// Copyright (c) Google LLC 2019
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file or at
// https://opensource.org/licenses/MIT.

#include "./context.h"

#include "./platform.h"
#include <brunsli/types.h>

namespace brunsli {

static const double kSqrt2 = 1.414213562;
static const int kSqrt2FixedPoint =
    static_cast<int>(kSqrt2 * kACPredictPrecision);

void ComputeACPredictMultipliers(const int* quant,
                                 int* mult_row,
                                 int* mult_col) {
  for (size_t y = 0; y < 8; ++y) {
    for (size_t x = 0; x < 8; ++x) {
      mult_row[x + 8 * y] =
          (quant[x + 8 * y] * kSqrt2FixedPoint) / quant[y * 8];
      // mult_col is transposed i.e. destination rows and columns are swapped.
      mult_col[x * 8 + y] = (quant[x + 8 * y] * kSqrt2FixedPoint) / quant[x];
    }
  }
}

// ComponentStateDC

void ComponentStateDC::InitAll() {
  is_zero_prob.Init(135);
  for (size_t i = 0; i < sign_prob.size(); ++i) {
    sign_prob[i].Init(128);
  }
  for (size_t i = 0; i < is_empty_block_prob.size(); ++i) {
    is_empty_block_prob[i].Init(74);
  }
  for (size_t i = 0; i < first_extra_bit_prob.size(); ++i) {
    first_extra_bit_prob[i].Init(150);
  }
}

// ComponentState

static const uint8_t kInitProb[64] = {
  // First probability (dc) and last of this table are unused.
  228, 216, 216, 195, 192, 189, 182, 184,
  179, 176, 171, 168, 166, 159, 156, 151,
  151, 150, 150, 146, 144, 138, 138, 137,
  135, 131, 127, 126, 124, 123, 124, 123,
  122, 121, 118, 117, 114, 115, 116, 116,
  115, 115, 114, 111, 111, 111, 112, 111,
  110, 110, 110, 111, 111, 114, 110, 111,
  112, 113, 116, 120, 126, 131, 147, 160,
};

static const uint8_t
    kInitProbNonzero[kNumNonZeroContextCount][kNumNonZeroTreeSize] = {
        {251, 252, 117, 249, 161, 136, 83,  238, 184, 126, 137, 129, 140,
         119, 70,  213, 160, 175, 174, 130, 166, 134, 122, 125, 131, 144,
         136, 133, 139, 123, 79,  216, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {254, 252, 174, 232, 189, 155, 122, 177, 204, 173, 146, 149, 141,
         133, 103, 109, 167, 187, 168, 142, 154, 147, 125, 139, 144, 138,
         138, 153, 141, 133, 90,  121, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {251, 240, 197, 176, 184, 177, 114, 89,  194, 165, 153, 161, 158,
         136, 92,  95,  123, 171, 160, 140, 148, 136, 129, 139, 145, 136,
         143, 134, 138, 124, 92,  154, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {247, 220, 201, 110, 194, 176, 147, 59,  175, 171, 156, 157, 152,
         146, 115, 114, 88,  151, 164, 141, 153, 135, 141, 131, 146, 139,
         140, 145, 138, 137, 112, 184, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {238, 179, 203, 63,  194, 173, 149, 71,  139, 169, 154, 159, 150,
         146, 117, 143, 78,  122, 152, 137, 149, 138, 138, 133, 134, 142,
         142, 142, 148, 128, 118, 199, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {227, 127, 200, 44,  192, 170, 148, 100, 102, 161, 156, 153, 148,
         149, 124, 160, 88,  101, 134, 132, 149, 145, 134, 134, 136, 141,
         138, 142, 144, 137, 116, 208, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {214, 86,  195, 44,  187, 163, 148, 126, 81,  147, 156, 152, 150,
         144, 121, 172, 96,  95,  117, 122, 145, 152, 136, 133, 135, 135,
         131, 142, 141, 135, 114, 217, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {198, 56,  191, 54,  171, 162, 147, 144, 74,  128, 152, 149, 150,
         142, 119, 177, 101, 100, 106, 111, 135, 154, 136, 137, 136, 132,
         133, 142, 144, 130, 117, 222, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {176, 40,  189, 73,  147, 159, 148, 152, 79,  106, 147, 149, 151,
         139, 123, 188, 108, 110, 106, 97,  125, 151, 137, 138, 135, 135,
         134, 136, 140, 131, 116, 221, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {148, 33,  185, 88,  117, 158, 145, 163, 95,  91,  137, 146, 150,
         140, 120, 197, 115, 116, 114, 92,  114, 144, 130, 133, 132, 133,
         129, 140, 138, 130, 111, 224, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {117, 31,  180, 104, 93,  150, 143, 166, 99,  85,  124, 139, 148,
         142, 118, 201, 105, 120, 120, 90,  107, 135, 127, 130, 131, 131,
         132, 140, 142, 133, 114, 229, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {87,  35,  170, 110, 78,  141, 144, 176, 106, 90,  112, 132, 143,
         138, 119, 204, 111, 121, 125, 90,  105, 131, 124, 122, 129, 128,
         129, 137, 138, 133, 114, 227, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {63,  42,  159, 123, 73,  127, 142, 191, 105, 91,  105, 123, 139,
         137, 120, 209, 117, 110, 122, 98,  110, 125, 115, 123, 122, 126,
         128, 134, 141, 129, 113, 229, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {45,  53,  146, 135, 71,  114, 138, 193, 100, 98,  98,  113, 133,
         135, 118, 222, 113, 111, 139, 103, 107, 126, 111, 119, 121, 122,
         127, 135, 141, 128, 114, 242, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {33,  60,  132, 138, 75,  100, 134, 203, 112, 99,  98,  105, 126,
         131, 115, 229, 107, 93,  121, 106, 108, 122, 106, 109, 114, 116,
         127, 133, 143, 128, 110, 242, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {24,  70,  118, 134, 76,  87,  130, 201, 110, 96,  99,  97,  119,
         130, 111, 229, 97,  104, 125, 102, 112, 125, 101, 109, 113, 114,
         125, 129, 142, 127, 112, 241, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {17,  65,  100, 121, 80,  75,  124, 174, 117, 100, 94,  93,  114,
         128, 110, 216, 103, 94,  113, 122, 118, 126, 113, 108, 105, 108,
         122, 128, 141, 125, 113, 238, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {12,  70,  82,  132, 78,  65,  118, 155, 136, 103, 97,  89,  106,
         124, 111, 215, 115, 123, 129, 99,  104, 127, 110, 108, 101, 109,
         118, 126, 136, 123, 110, 233, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {8,   66,  61,  117, 91,  59,  108, 195, 101, 112, 99,  99,  99,
         116, 106, 230, 127, 99,  144, 101, 118, 137, 117, 111, 106, 104,
         116, 121, 134, 122, 110, 223, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {6,   78,  42,  146, 101, 54,  94,  201, 116, 102, 110, 94,  92,
         108, 103, 214, 108, 111, 127, 102, 121, 132, 120, 121, 95,  98,
         110, 121, 129, 117, 107, 235, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {5,   93,  29,  145, 102, 52,  77,  216, 108, 115, 108, 102, 89,
         97,  94,  229, 89,  103, 139, 120, 103, 151, 102, 100, 97,  96,
         99,  111, 125, 116, 104, 242, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {4,   105, 21,  145, 100, 54,  64,  217, 100, 122, 128, 87,  88,
         91,  87,  230, 112, 80,  148, 95,  146, 123, 96,  140, 90,  91,
         98,  106, 122, 111, 100, 249, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {4,   130, 14,  142, 104, 56,  51,  208, 116, 135, 100, 89,  82,
         84,  75,  239, 85,  85,  122, 125, 94,  144, 151, 136, 92,  97,
         104, 109, 113, 110, 91,  246, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {3,   126, 9,   172, 105, 57,  39,  219, 95,  120, 118, 96,  93,
         75,  66,  241, 102, 134, 96,  156, 146, 162, 130, 112, 82,  89,
         97,  101, 116, 103, 82,  254, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {3,   149, 7,   182, 122, 54,  29,  224, 103, 100, 113, 96,  90,
         74,  55,  250, 127, 94,  118, 93,  135, 160, 113, 130, 95,  117,
         106, 96,  111, 97,  77,  242, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {3,   150, 4,   170, 138, 59,  20,  229, 91,  150, 107, 98,  92,
         68,  48,  245, 113, 64,  114, 111, 134, 127, 102, 104, 85,  118,
         103, 107, 102, 91,  72,  245, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {3,   171, 3,   165, 137, 62,  14,  211, 96,  127, 132, 121, 95,
         62,  37,  248, 102, 57,  144, 85,  127, 191, 102, 97,  127, 104,
         91,  102, 107, 81,  64,  254, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {2,   166, 2,   196, 122, 65,  10,  243, 102, 93,  117, 92,  96,
         63,  29,  251, 169, 159, 149, 96,  91,  139, 157, 40,  100, 89,
         120, 92,  109, 79,  58,  247, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {2,   176, 2,   189, 118, 48,  7,   219, 68,  43,  109, 96,  129,
         75,  19,  254, 2,   3,   185, 6,   102, 127, 127, 127, 1,   131,
         83,  99,  107, 80,  45,  254, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {1,   205, 2,   208, 64,  89,  4,   223, 29,  169, 29,  123, 118,
         76,  11,  240, 202, 243, 65,  6,   12,  243, 96,  55,  102, 102,
         114, 102, 107, 74,  31,  247, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {1,   216, 1,   214, 127, 94,  2,   234, 145, 3,   127, 106, 155,
         80,  4,   247, 4,   65,  86,  127, 127, 127, 127, 102, 127, 143,
         143, 108, 113, 80,  16,  216, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128},
        {2,   199, 1,   222, 93,  94,  1,   232, 2,   65,  74,  139, 201,
         48,  2,   254, 169, 127, 52,  243, 251, 249, 102, 86,  202, 153,
         65,  65,  146, 69,  8,   238, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128,
         128, 128, 128, 128, 128, 128, 128, 128, 128, 128, 128}};

void ComponentState::InitAll() {
  for (int i = 0; i < kNumNonzeroBuckets; ++i) {
    for (int k = 0; k < kDCTBlockSize; ++k) {
      const int v = kInitProb[k] + 9 * (i - 7);
      BRUNSLI_DCHECK(v <= 255);
      is_zero_prob[i * kDCTBlockSize + k].Init(v);
    }
  }
  for (size_t i = 0; i < sign_prob.size(); ++i) {
    if (i < kMaxAverageContext * kDCTBlockSize) {
      sign_prob[i].Init(108);
    } else if (i < (kMaxAverageContext + 1) * kDCTBlockSize) {
      sign_prob[i].Init(128);
    } else {
      sign_prob[i].Init(148);
    }
  }
  for (size_t i = 0; i < first_extra_bit_prob.size(); ++i) {
    first_extra_bit_prob[i].Init(158);
  }

  for (size_t i = 0; i < kNumNonZeroContextCount; ++i) {
    Prob* non_zero_probs = num_nonzero_prob + i * kNumNonZeroTreeSize;
    for (size_t j = 0; j < kNumNonZeroTreeSize; ++j) {
      non_zero_probs[j].Init(kInitProbNonzero[i][j]);
    }
  }
}

}  // namespace brunsli
